#!/usr/bin/perl
#
# express-bundle-pool.pl
# Create a union of bundle_ids from multiple eXpress 'results.xprs' files.
# 20140221
#
# Copyright (C) 2014 Shawn Driscoll, sdriscoll@salk.edu
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#
# Description:
# Pools bundle_id to transcript ids from multiple express runs to create a 
# final bundle set. Between files intersections of bundles are found and those
# bundles that intersect are merged.  eXpress can reveal interesting connections
# between genes via the bundle_id column however it produces them on a sample-to-sample
# bases and provides no way to create a union of the bundles to encompass an
# entire project.
#

use strict;
use Getopt::Std;

# -----------------------------------------------------------------------------
# parse options
# -----------------------------------------------------------------------------

my %opts = (
	"h" => 0,
);
my $r = getopt("h", \%opts);

if(!$r || $opts{"h"}) {
	die usage();
}

# -----------------------------------------------------------------------------
# check args
# -----------------------------------------------------------------------------

die "Error: too few arguments!\n\n" . usage() unless (scalar @ARGV) >= 2;
my (@files) = @ARGV;

# -----------------------------------------------------------------------------
# setup vars
# -----------------------------------------------------------------------------

my $num_files = scalar @files;
my $fname;
my $tmp;
my @aln;
my $res;
my (@bundles, @bundles_tmp);
my %map;
my ($i, $j, $k) = (0, 0, 0);
my $b;
my $merged = 0;
my $bid;
my (@idx, %idx_set);

# -----------------------------------------------------------------------------
# main script
# -----------------------------------------------------------------------------

# load the first file
$res = load_file(\@bundles, $files[0]);
print STDERR "# loaded $files[0]\n";

# load each additional file and merge into this first set
for($i=1; $i<$num_files; $i++) {
	print STDERR "# updating bundle map\n";
	$res = make_bundle_map(\%map, \@bundles);
	print STDERR "# loading $files[$i]\n";
	@bundle_tmp = ();
	$res = load_file(\@bundle_tmp, $files[$i]);
	# loop through and try to merge with the first one
	print STDERR "# merging...\n";
	foreach $b (@bundle_tmp) {
		# loop through bundles in the new file

		# for each key in %$b find the current bundle id from map
		@idx = ();
		foreach (keys %$b) {
			push(@idx, $map{$_});
		}
		# reduce array of bundle ids down to a unique set
		%idx_set = ();
		ar_to_set(\%idx_set, \@idx);
		@idx = keys %idx_set;
		# if there is more than one bundle id in the returned set then two or 
		# more of the former bundles now overlap due to the new information
		# in the current bundle. We'll merge each of the sets in @idx into
		# the first index in @idx ($idx[0]) then destroy the redundant sets.
		if((scalar @idx) > 1) {
			# need to merge some sets
			for($k=1; $k<(scalar @idx); $k++) {
				# create union into first index
				$res = set_union(\%{$bundles[$idx[0]]}, \%{$bundles[$idx[$k]]});
				# remove redundant set
				$bundles[$idx[$k]] = 0;
			}
			# because the set list was revised we have to update the bundle map
			$res = make_bundle_map(\%map, \@bundles);
		}
	}
}

# print bundles out
$i = 0;
foreach $b (@bundles) {
	if($b) {
		$i++;
		$bid = sprintf("BID%0.8d", $i);
		foreach (keys %$b) {
			print STDOUT join("\t", $bid, $_) . "\n";
		}		
	}
}

# -----------------------------------------------------------------------------
# subs
# -----------------------------------------------------------------------------

# --
# load the bundle id to transcript id table from a single .xprs file. The 
# bundle array is passed by reference (first parameter)
# load_file(\@a, $file_name)
sub load_file {
	my ($ar_bundles, $fname) = @_;
	my @tmp;
	my @aln;
	my %b;
	my %bundles = ();

	# make sure the bundle array is empty
	@$ar_bundles = ();

	open IN, '<', $fname || return 1;
	# pull the transcript ids into a hash keyed by the bundle_id column
	while(<IN>) {
		next if /^bundle_id/;
		chomp;	
		@aln = split(/\t/);

		if(! exists $bundles{$aln[0]}) {
			my @tmp = ();
			$bundles{$aln[0]} = \@tmp;
		}

		push(@{$bundles{$aln[0]}}, $aln[1]);
	}

	close IN;

	# now the hash of array of transcripts has to be converted into an 
	# array of hashes.

	foreach (keys %bundles) {
		my %b = ();
		ar_to_set(\%b, \@{$bundles{$_}});
		push(@$ar_bundles, \%b);
	}

	return 0;
}

# --
# make_bundle_map(\%map, \@bundles)
# build a hash of transcript ids to bundle ids. the bundle ids are just the
# array indices of the bundle in @bundles.
sub make_bundle_map {
	my ($map, $bundles) = @_;

	# bundles is an array of hashes, map is an empty hash
	%$map = ();

	my $i = scalar @$bundles;
	my $b;

	while($i--) {
		# sum bundles may have been dropped if any merging has happened
		if($$bundles[$i]) {
			foreach (keys %{$$bundles[$i]}) {
				$$map{$_} = $i;
			}
		}
	}

	return 0;
}

# --
# ar_to_set(\%set_out, \@ar_in)
# converts array of values into a set (a hash)
sub ar_to_set {

	my ($set_out, $ar_in) = @_;

	foreach (@$ar_in) {
		if(! exists $$set_out{$_}) {
			$$set_out{$_} = 0;
		}
	}

	return 0;
}

# --
# set_union(\%set1, \%set2)
# create union of two sets into set1
sub set_union {
	my ($set1, $set2) = @_;

	foreach (keys %$set2) {
		$$set1{$_} = 0 if ! exists $$set1{$_};
	}

	return 0;
}


sub usage {
	return "Usage: express-bundle-pool.pl [options] xprs1 xprs2 [additional xprs outputs]

Create unions of bundle_id pools from multiple eXpress runs.

Options:
  -h             Display this message and exit

";
}
